#-*- coding:utf8 -*-
import socket
import select
import os
import errno
import Queue
import json
import threading
import time

class Logger:
    m_loggerLock = threading.Lock()

    def __init__(self):
        return
    
    @staticmethod
    def writeLoger(data):
        Logger.m_loggerLock.acquire()
        print "[%s]> %s"%(time.asctime(time.localtime()),data)
        Logger.m_loggerLock.release()

# {"head":{"type":"service"},"body":{"service":"xxx"}}
# {"head":{"type":"manage"},"body":{"type":"register","service":"xxx"}}
# {"head":{"type":"service","fromsock":1},"body":{"service":"xxx"}}
# {"head":{"type":"manage","fromsock":1},"body":{"type":"register","service":"xxx"}}

class MsgDecode:
    def __init__(self,data):
        self.m_data = data
        self.m_st_data = json.loads(self.m_data)
    
    def getHeadMsgType(self):
        return self.m_st_data["head"]["type"]
        
    def getFromSock(self):
        return self.m_st_data["head"]["fromsock"]
        
    def getBodyMsgType(self):
        return self.m_st_data["body"]["type"]
        
    def isManageMsg(self):
        return self.getHeadMsgType() == "manage"
        
    def getMgrMsgType(self):
        return self.getBodyMsgType()

    def isRegisterMgrMsg(self):
        return self.getMgrMsgType() == "register"
    
    def getRegisterService(self):
        return self.m_st_data["body"]["service"]
    
    def getService(self):
        return self.m_st_data["body"]["service"].encode('ascii')
        
class MsgDispatcher:
    def __init__(self,inaddr,inport,outaddr,outport):
        # 输入tcp端口
        self.m_inaddr = inaddr
        self.m_inport = inport
        self.m_insock = socket.socket(socket.AF_INET,socket.SOCK_STREAM)
        # 分发tcp端口
        self.m_outaddr = outaddr
        self.m_outport = outport
        self.m_outsock = socket.socket(socket.AF_INET,socket.SOCK_STREAM)
        # 输入tcp的客户端映射
        self.m_insock_map = {} # fd -> sock 
        self.m_inaddr_map = {} # fd -> addr
        # 分发tcp的客户端映射
        self.m_outsock_map = {} # fd -> sock
        self.m_outaddr_map = {} # fd -> addr
        # 输入消息缓冲队列
        self.m_inqueue = Queue.Queue()
        # 响应消息缓冲队列
        self.m_outqueue = Queue.Queue()
        # 管理消息缓冲队列
        self.m_mgrqueue = Queue.Queue()
        # 分发服务注册列表映射
        self.m_servicemap = {} # service -> fdlst
        self.m_serviceidx = {} # service -> clientIdx


    def __del__(self):
        return

    def initDispatcher(self):
        # 启动监听资源
        Logger.writeLoger("Start in at %s:%d"%(self.m_inaddr,self.m_inport))
        Logger.writeLoger("Start out at %s:%d"%(self.m_outaddr,self.m_outport))
        self.m_insock.bind((self.m_inaddr,self.m_inport))
        self.m_insock.listen(5)
        self.m_outsock.bind((self.m_outaddr,self.m_outport))
        self.m_outsock.listen(5)

    def openInSock(self,client,address):
        Logger.writeLoger("[IN] %s:%d(fd:%d) connect to in sock."%(address[0],address[1],client.fileno()))
        client.setblocking(0)
        # 注册输入端映射
        self.m_insock_map[client.fileno()] = client
        self.m_inaddr_map[client.fileno()] = address    

    def closeInSock(self,client):
        Logger.writeLoger("[IN] %s:%d closed" % (self.m_inaddr_map[client.fileno()][0], self.m_inaddr_map[client.fileno()][1]))
        # 移除输入端映射
        self.m_insock_map.pop(client.fileno())
        self.m_inaddr_map.pop(client.fileno())
        client.close()

    def openOutSock(self,client,address):
        Logger.writeLoger("[OUT] %s:%d(fd:%d) connect to in sock."%(address[0],address[1],client.fileno()))
        client.setblocking(0)
        # 注册输入端映射
        self.m_outsock_map[client.fileno()] = client
        self.m_outaddr_map[client.fileno()] = address    

    def closeOutSock(self,client):
        Logger.writeLoger("[OUT] %s:%d closed" % (self.m_inaddr_map[client.fileno()][0], self.m_inaddr_map[client.fileno()][1]))
        # 移除输入端映射
        self.m_outsock_map.pop(client.fileno())
        self.m_outaddr_map.pop(client.fileno())
        client.close()

    def sendDataInSock(self,fileno,data):
        Logger.writeLoger("[IN] send to %d data:%s"%(fileno,data))
        return self.m_insock_map[fileno].send(data) 

    def sendDataOutSock(self,fileno,data):
        Logger.writeLoger("[OUT] send to %d data:%s"%(fileno,data))
        return self.m_outsock_map[fileno].send(data) 

    def recvDataInSock(self,fileno):
        # 有 可读 事件激活
        datas = ''
        while True:
            try:
                # 从激活 fd 上 recv 10 字节数据
                data = self.m_insock_map[fileno].recv(10)
                # 若当前没有接收到数据，并且之前的累计数据也没有
                if not data and not datas:
                # 从 epoll 句柄中移除该 连接 fd
                    self.closeInSock(self.m_insock_map[fileno])
                    break
                else:
                # 将接收到的数据拼接保存在 datas 中
                    datas += data
            except socket.error, msg:
                # 在 非阻塞 socket 上进行 recv 需要处理 读穿 的情况
                # 这里实际上是利用 读穿 出 异常 的方式跳到这里进行后续处理
                if msg.errno == errno.EAGAIN or msg.errno == 10035:
                    Logger.writeLoger("[IN] %s:%d(fd:%s) receive %s" % (self.m_inaddr_map[fileno][0],self.m_inaddr_map[fileno][1],fileno, datas))
                    # 将已接收数据保存起来
                    break
                else:
                    # 出错处理
                    Logger.writeLoger(msg)
                    self.closeInSock(self.m_insock_map[fileno])
                    break
        return datas

    def recvDataOutSock(self,fileno):
        # 有 可读 事件激活
        datas = ''
        while True:
            try:
                # 从激活 fd 上 recv 10 字节数据
                data = self.m_outsock_map[fileno].recv(10)
                # 若当前没有接收到数据，并且之前的累计数据也没有
                if not data and not datas:
                # 从 epoll 句柄中移除该 连接 fd
                    self.closeOutSock(self.m_outsock_map[fileno])
                    break
                else:
                # 将接收到的数据拼接保存在 datas 中
                    datas += data
            except socket.error, msg:
                # 在 非阻塞 socket 上进行 recv 需要处理 读穿 的情况
                # 这里实际上是利用 读穿 出 异常 的方式跳到这里进行后续处理
                if msg.errno == errno.EAGAIN or msg.errno == 10035:
                    Logger.writeLoger("[OUT] %s:%d(fd:%s) receive %s" % (self.m_outaddr_map[fileno][0],self.m_outaddr_map[fileno][1],fileno, datas))
                    # 将已接收数据保存起来
                    break
                else:
                    # 出错处理
                    Logger.writeLoger(msg)
                    self.closeOutSock(self.m_outsock_map[fileno])
                    break
        return datas

    def procInSock(self,timeout):
        return

    def procOutSock(self,timeout):
        return

    def reencodeMsg(self,data,fileno):
        stdata = json.loads(data)
        if not stdata["head"].has_key("fromsock"):
            stdata["head"]["fromsock"] = fileno
        #print stdata
        return json.dumps(stdata,ensure_ascii=False)
    
    def pushIntoInQueue(self,data,fileno):
        data=self.reencodeMsg(data,fileno)
        #print data
        msgdecode = MsgDecode(data)
        if msgdecode.isManageMsg():
            Logger.writeLoger("[IN] in queue recv a mgr msg, push it into manage queue:%s"%(data))
            self.pushIntoMgrQueue(data)
            return 
        self.m_inqueue.put(data)

    def popFromInQueue(self,timeout1):
        try:
            data = self.m_inqueue.get(timeout=timeout1)
            return data
        except Exception, e:
            #print e
            return ""

    def pushIntoOutQueue(self,data,fileno):
        data = self.reencodeMsg(data,fileno)
        msgdecode = MsgDecode(data)
        if msgdecode.isManageMsg():
            Logger.writeLoger("[OUT] out queue recv a mgr msg, push it into manage queue:%s"%(data))
            self.pushIntoMgrQueue(data)
            return 
        self.m_outqueue.put(data)

    def popFromOutQueue(self,timeout1):
        try:
            data = self.m_outqueue.get(timeout=timeout1)
            return data
        except Exception, e:
            #print e
            return ""

    def pushIntoMgrQueue(self,data):
        self.m_mgrqueue.put(data)

    def popFromMgrQueue(self,timeout1):
        try:
            data = self.m_mgrqueue.get(timeout=timeout1)
            return data
        except Exception, e:
            #print e
            return ""
            
    def dispatchQueueMsg(self):
        data = self.popFromInQueue(10)
        if len(data) == 0:
            return
        msgdecode = MsgDecode(data)
        servicename = msgdecode.getService()
        if not self.m_servicemap.has_key(servicename):
            Logger.writeLoger("[%s] no service [%s] registered."%(str(threading.current_thread().getName()),servicename))
        else:
            if self.m_serviceidx[servicename] == len(self.m_servicemap[servicename]):
                self.m_serviceidx[servicename] = 0
            idx = self.m_serviceidx[servicename]
            Logger.writeLoger("[%s] forward service [%s] to fd [%d]."%(str(threading.current_thread().getName()),servicename,self.m_servicemap[servicename][idx]))
            self.sendDataOutSock(self.m_servicemap[servicename][idx],data)
            self.m_serviceidx[servicename]=self.m_serviceidx[servicename]+1

    def mgrMsgProcess(self):
        data = self.popFromMgrQueue(10)
        if len(data) == 0:
            return
        msgdecode = MsgDecode(data)
        if msgdecode.isRegisterMgrMsg():
            Logger.writeLoger("[MGR] socket %d register service %s"%(msgdecode.getFromSock(),msgdecode.getRegisterService()))
            if not self.m_servicemap.has_key(msgdecode.getRegisterService()):
                self.m_servicemap[msgdecode.getRegisterService()] = []
                self.m_serviceidx[msgdecode.getRegisterService()] = 0
            self.m_servicemap[msgdecode.getRegisterService()].append(msgdecode.getFromSock())
            Logger.writeLoger("[MGR] service %s clientlist:%s"%(msgdecode.getRegisterService(),str(self.m_servicemap[msgdecode.getRegisterService()])))
        else:
            Logger.writeLoger("unkown mgr msg:%s"%(msgdecode.getMgrMsgType()))

class EpollMsgDispatcher(MsgDispatcher):
    def __init__(self,inaddr,inport,outaddr,outport):
        MsgDispatcher.__init__(self,inaddr,inport,outaddr,outport)
        self.m_inepoll = select.epoll()
        self.m_outepoll = select.epoll()
        
    def initDispatcher(self):
        MsgDispatcher.initDispatcher(self)
        self.m_inepoll.register(self.m_insock.fileno(),select.EPOLLIN)
        self.m_outepoll.register(self.m_outsock.fileno(),select.EPOLLIN)

    def openInSock(self,client,address):
        MsgDispatcher.openInSock(self,client,address)
        self.m_inepoll.register(client.fileno(),select.EPOLLIN | select.EPOLLET)

    def closeInSock(self,client):
        MsgDispatcher.closeInSock(self,client)
        self.m_inepoll.unregister(client.fileno())

    def procInSock(self,timeout):
        events = self.m_inepoll.poll(timeout)
        for fileno,event in events:
            if fileno == self.m_insock.fileno():
                (client,address) = self.m_insock.accept()
                #print "%s:%d connect to in sock."%(address[0],address[1])
                #将新的连接进行注册，用来接收消息
                openInSock(client,address)
            elif select.EPOLLIN & events:
                recvDataInSock(fileno)
            elif select.EPOLLHUP & events:
                closeInSock(fileno)
            elif select.EPOLLOUT & events:
                sendDataInSock(fileno,"")
            else:
                # 其他 epoll 事件不进行处理
                continue

class SelectMsgDispatcher(MsgDispatcher):
    def __init__(self,inaddr,inport,outaddr,outport):
        MsgDispatcher.__init__(self,inaddr,inport,outaddr,outport)
        # in端口的列表
        self.m_select_readlist = []
        self.m_select_writelist = []
        self.m_select_exceptionlist = []
        # out端口的列表
        self.m_out_select_readlist = []
        self.m_out_select_writelist = []
        self.m_out_select_exceptionlist = []
        
    def initDispatcher(self):
        MsgDispatcher.initDispatcher(self)
        self.m_select_readlist.append(self.m_insock)
        self.m_select_writelist.append(self.m_insock)
        self.m_select_exceptionlist.append(self.m_insock)

        self.m_out_select_readlist.append(self.m_outsock)
        self.m_out_select_writelist.append(self.m_outsock)
        self.m_out_select_exceptionlist.append(self.m_outsock)

    def openInSock(self,client,address):
        MsgDispatcher.openInSock(self,client,address)
        self.m_select_readlist.append(client)
        self.m_select_writelist.append(client)
        self.m_select_exceptionlist.append(client)

    def closeInSock(self,client):
        #MsgDispatcher.closeInSock(self,client)
        Logger.writeLoger("[IN] %s:%d closed" % (self.m_inaddr_map[client.fileno()][0], self.m_inaddr_map[client.fileno()][1]))
        self.m_select_readlist.remove(client)
        self.m_select_writelist.remove(client)
        self.m_select_exceptionlist.remove(client)

    def openOutSock(self,client,address):
        MsgDispatcher.openOutSock(self,client,address)
        self.m_out_select_readlist.append(client)
        self.m_out_select_writelist.append(client)
        self.m_out_select_exceptionlist.append(client)

    def closeOutSock(self,client):
        #MsgDispatcher.closeInSock(self,client)
        Logger.writeLoger("[OUT] %s:%d closed" % (self.m_outaddr_map[client.fileno()][0], self.m_outaddr_map[client.fileno()][1]))
        self.m_out_select_readlist.remove(client)
        self.m_out_select_writelist.remove(client)
        self.m_out_select_exceptionlist.remove(client)

    # 接收新消息并放到Queue中去
    def procInSock(self,timeout):
        rs,ws,es=select.select(self.m_select_readlist,self.m_select_writelist,self.m_select_exceptionlist)
        #if len(rs) > 0 or len(ws) > 0 or len(es) > 0:
        #    print rs
        #    print ws 
        #    print es
        for r in rs:
            if r is self.m_insock:
                (client,address) = self.m_insock.accept()
                #print "%s:%d connect to in sock."%(address[0],address[1])
                #将新的连接进行注册，用来接收消息
                self.openInSock(client,address)
            else:
                data = self.recvDataInSock(r.fileno())
                if data == "":
                    self.closeOutSock(r)
                else:
                    self.pushIntoInQueue(data,r.fileno())
        #for w in ws:
        #    self.sendDataInSock(w.fileno(),"")
        for e in es:
            self.closeInSock(e)

    # 接收新消息并放到Queue中去
    def procOutSock(self,timeout):
        rs,ws,es=select.select(self.m_out_select_readlist,self.m_out_select_writelist,self.m_out_select_exceptionlist)
        #if len(rs) > 0 or len(ws) > 0 or len(es) > 0:
        #    print rs
        #    print ws 
        #    print es
        for r in rs:
            if r is self.m_outsock:
                (client,address) = self.m_outsock.accept()
                #print "%s:%d connect to in sock."%(address[0],address[1])
                #将新的连接进行注册，用来接收消息
                self.openOutSock(client,address)
            else:
                data = self.recvDataOutSock(r.fileno())
                if data == "":
                    self.closeOutSock(r)
                else:
                    self.pushIntoOutQueue(data,r.fileno())
        #for w in ws:
            #self.sendDataOutSock(w.fileno(),"")
        for e in es:
            self.closeOutSock(e)

def oneDispatchWorker(msgdispatcher):
    #print threading.current_thread().getName()
    while True:
        msgdispatcher.dispatchQueueMsg()
        #time.sleep(1)

def oneInRecvSendWorker(msgdispatcher):
    while True:
        msgdispatcher.procInSock(10)

def oneOutRecvSendWorker(msgdispatcher):
    while True:
        msgdispatcher.procOutSock(10)

def oneMgrMsgProcWorker(msgdispatcher):
    while True:
        msgdispatcher.mgrMsgProcess()
        
def main():
    msgdis = SelectMsgDispatcher("127.0.0.1",1001,"127.0.0.1",1002)
    msgdis.initDispatcher()
    workers = []
    for i in range(1):
        w = threading.Thread(target=oneDispatchWorker, args=(msgdis,))
        workers.append(w)
    w = threading.Thread(target=oneInRecvSendWorker, args=(msgdis,))
    workers.append(w)
    w = threading.Thread(target=oneOutRecvSendWorker, args=(msgdis,))
    workers.append(w)
    w = threading.Thread(target=oneMgrMsgProcWorker, args=(msgdis,))
    workers.append(w)

    for w in workers:
        w.start()

    for w in workers:
        w.join()

if __name__ == '__main__':
    main()
